package network

import (
	"github.com/sirupsen/logrus"
	"net"
	"sync"
)

type TcpServer struct {
	Addr         string         //监听地址
	MaxConnNum   int            //最大连接数
	listener     net.Listener   //listener对象
	wgListener   sync.WaitGroup //关闭监听的WaitGroup对象
	mutexConnSet sync.Mutex     //判断连接数的锁
	connSet      ConnSet        //判断连接数的map
}

// Start 启动tcp服务器
func (server *TcpServer) Start() {
	server.init()
	go server.run()
}

func (server *TcpServer) init() {
	listener, err := net.Listen("tcp", server.Addr)
	if err != nil {
		logrus.Info(err)
	}
	server.listener = listener
	logrus.Infof("tcp server listen %s", server.Addr)

}

func (server *TcpServer) run() {
	server.wgListener.Add(1)
	defer server.wgListener.Done()
	for {
		conn, err := server.listener.Accept()
		if err != nil {
			logrus.Error(err)
			return
		}
		server.mutexConnSet.Lock()
		//检查最大连接数
		if len(server.connSet) >= server.MaxConnNum {
			server.mutexConnSet.Unlock()
			conn.Close()
			logrus.Debugf("too many connections! num： %d", server.MaxConnNum)
			return
		}
		server.connSet[conn] = struct{}{}
		server.mutexConnSet.Unlock()
	}
}

// Close 关闭tcp服务器
func (server *TcpServer) Close() {
	server.listener.Close()
	server.wgListener.Wait()
	server.mutexConnSet.Lock()
	for conn := range server.connSet {
		conn.Close()
	}
	server.connSet = nil
	server.mutexConnSet.Unlock()
	logrus.Info("tcp server closed...")
}
